package main

import (
	"context"
	"fmt"
	"github.com/IBM/sarama"
	"github.com/gin-gonic/gin"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	brokers     = []string{"localhost:29092", "localhost:39092", "localhost:49092"}
	messageChan = make(chan string, 100)
	topic       = "sse-demo"
)

func main() {

	// 初始化生产者
	producer, err := initProducer()
	if err != nil {
		panic(err)
	}

	consumer, err := initConsumer()
	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup
	wg.Add(1)
	go consumeMsg(ctx, consumer, &wg)

	r := gin.Default()

	r.POST("/send", sendHandler(producer))
	r.GET("/stream", sseHandler)

	// 启动服务
	go func() {
		if err := r.Run(":8080"); err != nil {
			fmt.Println("server error")
		}
	}()

	// 优雅关闭
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	cancel()
	wg.Wait()
	close(messageChan)
	producer.Close()
	consumer.Close()

}

// 消费者组处理器

// ConsumerGroupHandler 实现 sarama.ConsumerGroupHandler 接口
type ConsumerGroupHandler struct {
}

func (h ConsumerGroupHandler) Setup(_ sarama.ConsumerGroupSession) error {
	return nil
}
func (h ConsumerGroupHandler) Cleanup(_ sarama.ConsumerGroupSession) error { return nil }

func (h ConsumerGroupHandler) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	for msg := range claim.Messages() {
		fmt.Printf("Received message: Topic=%s Partition=%d Offset=%d Key=%s Value=%s\n",
			msg.Topic, msg.Partition, msg.Offset, string(msg.Key), string(msg.Value))
		sess.MarkMessage(msg, "")
	}
	return nil
}

func initProducer() (sarama.SyncProducer, error) {
	config := sarama.NewConfig()
	config.Producer.RequiredAcks = sarama.WaitForAll // 等待所有副本确认
	config.Producer.Retry.Max = 5                    // Retry up to 5 times to produce the message
	config.Producer.Return.Successes = true
	config.Version = sarama.V4_0_0_0
	return sarama.NewSyncProducer(brokers, config)
}
func sendHandler(producer sarama.SyncProducer) gin.HandlerFunc {
	return func(c *gin.Context) {
		msg := c.PostForm("msg")
		if msg == "" {
			c.JSON(400, gin.H{"error": "message is empty"})
			return
		}
		partition, offset, err := producer.SendMessage(&sarama.ProducerMessage{
			Topic: topic,
			Value: sarama.StringEncoder(msg),
		})
		fmt.Printf("Partition: %d, Offset: %d\n", partition, offset)
		if err != nil {
			c.JSON(500, gin.H{"error": "failed to send message"})
			return
		}
		c.JSON(200, gin.H{
			"status":  "msg sent success",
			"message": msg,
		})
	}
}

func consumeMsg(ctx context.Context, consumer sarama.ConsumerGroup, wg *sync.WaitGroup) {
	defer wg.Done()
	handler := ConsumerGroupHandler{}
	for {
		select {
		case <-ctx.Done():
			return
		default:
			err := consumer.Consume(ctx, []string{topic}, handler)
			if err != nil {
				fmt.Println("Error on consume err:", err)
				time.Sleep(1 * time.Second)
			}
		}
	}
}

func sseHandler(c *gin.Context) {
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")

	for {
		select {
		case msg := <-messageChan:
			fmt.Println(msg)
			c.Writer.Flush()
		case <-c.Writer.CloseNotify():
			return
		}
	}
}

func initConsumer() (sarama.ConsumerGroup, error) {
	config := sarama.NewConfig()
	config.Version = sarama.V1_0_0_0
	config.Consumer.Group.Rebalance.GroupStrategies = []sarama.BalanceStrategy{
		sarama.NewBalanceStrategyRange(),
	}
	config.Consumer.Offsets.Initial = sarama.OffsetOldest

	return sarama.NewConsumerGroup(brokers, "group1", config)
}
